from pymilvus import Collection, connections, db, FieldSchema, DataType, CollectionSchema, utility
from configs.load_configs import ConfigLoader

# 创建向量数据库Collection
# data 业务数据
from models.embedding.load_embedding_model import LoadEmbeddingModel


class MilvusDB(object):

    def __init__(self):
        config = ConfigLoader.load_milvus_config()
        # 连接到 Milvus 服务器
        connections.connect(host=config["host"], port=config["port"])

        self.model = LoadEmbeddingModel()
        self.model.load()

    def create_milvus(self, param_data):
        # 使用数据库
        db.using_database("default")

        # ============================================================

        # 创建主键
        pk_id = FieldSchema(
            name="pk_id",
            dtype=DataType.INT64,
            is_primary=True,
        )
        # 创建标量，用来记录DDL语句
        ddl = FieldSchema(
            name="ddl",
            dtype=DataType.VARCHAR,
            max_length=5000,
            default_value="Unknown"
        )
        # 创建向量，用来记录SQL语句
        sql = FieldSchema(
            name="sql",
            dtype=DataType.VARCHAR,
            max_length=5000,
            default_value="Unknown"
        )
        # 创建向量，用来记录SQL描述
        sql_describe = FieldSchema(
            name="sql_describe",
            dtype=DataType.VARCHAR,
            max_length=800,
            default_value="Unknown"
        )
        # 创建向量，用来记录SQL描述向量化
        sql_describe_vector = FieldSchema(
            name="sql_describe_vector",
            dtype=DataType.FLOAT_VECTOR,
            dim=1024
        )
        # 将列名放入CollectionSchema对象中，为创建Collection做准备
        schema = CollectionSchema(
            fields=[pk_id, ddl, sql, sql_describe, sql_describe_vector],
            description="根据旁白文字查询视频时间",
            enable_dynamic_field=True
        )
        # Collection名
        collection_name = "text2sql"

        # 在default数据库中创建Collection
        collection = Collection(
            name=collection_name,
            schema=schema,
            using='default'
        )

        # ============================================================

        # 创建向量索引
        index_params = {
            "metric_type": "L2",
            "index_type": "HNSW",
            "params": {
                "M": 8,
                "efConstruction": 64
            }
        }

        # 为Collection中的向量列名sql_describe_vector添加向量索引
        collection.create_index(
            field_name="sql_describe_vector",
            index_params=index_params
        )
        utility.index_building_progress("text2sql")

        # ============================================================

        # 创建一个新的列表来存储新的数据点
        new_data = []

        # 遍历每行数据
        for idx, row in enumerate(param_data):
            # 使用嵌入模型获取旁白文字的向量
            sql_describe = row['sql_describe']
            print(sql_describe)
            embedding_vector = self.model.encode(sql_describe)
            # 创建数据点字典
            data_point = {
                "pk_id": idx + 1,  # 主键
                "ddl": row['ddl'],  # 建表语句
                "sql": row['sql'],  # SQL语句
                "sql_describe": sql_describe,  # SQL描述
                "sql_describe_vector": embedding_vector.tolist()  # SQL描述向量
            }
            # 添加到数据列表
            new_data.append(data_point)

        # 插入数据到Milvus
        collection.insert(new_data)
        # 加载Collection
        collection.load()

    # 查询向量数据库
    # query_vector 查询条件向量
    def search_milvus(self, query_vector):
        # 将参数向量化
        query_vector = self.model.encode(query_vector)

        # 使用数据库
        db.using_database("default")

        # 获取现有的集合
        collection = Collection("text2sql")
        # 加载集合到内存
        collection.load()

        # 准备搜索参数
        search_params = {
            "metric_type": "L2",
            "offset": 0,
            "params": {"ef": 250}
        }

        # 进行搜索操作
        results = collection.search(
            data=[query_vector],  # 待搜索的数据，这里使用二维数组表示
            anns_field="sql_describe_vector",  # 指定进行向量搜索的列名
            # `param` 中 `offset` 和 `limit` 的总和应小于 16384
            param=search_params,  # 搜索参数，包括偏移和限制等
            limit=1,  # 搜索结果返回的数量限制
            expr=None,  # 表达式，用于进一步筛选结果，这里设置为 None
            # 指定要从搜索结果中检索的列名(非内置列名都要写)
            output_fields=['pk_id', 'ddl', 'sql', 'sql_describe'],
            consistency_level="Strong"  # 一致性级别，这里设置为“Strong”
        )

        # 输出结果格式化成横向表格
        result_data = [r.entity for r in results[0]]

        # 释放加载的集合，减少内存消耗
        # conn.release()

        # 断开与 Milvus 服务器的连接
        # conn.disconnect()

        return result_data


if __name__ == '__main__':
    db = MilvusDB()

    # 预置数据
    data = [
        {
            "ddl": "CREATE TABLE `contracts` (   `id` int(11) NOT NULL AUTO_INCREMENT COMMENT '合同ID',   `year` int(11) NOT NULL COMMENT '合同签订年份',   `department` varchar(255) NOT NULL COMMENT '部门名称,合同的所属部门',   `budget_contract_amount` decimal(10,2) NOT NULL COMMENT '年初预算的新签合同额',   `actual_contract_amount` decimal(10,2) NOT NULL COMMENT '实际的新签合同额',   `created_at` timestamp NOT NULL DEFAULT current_timestamp() COMMENT '创建时间',   `updated_at` timestamp NOT NULL DEFAULT current_timestamp() ON UPDATE current_timestamp() COMMENT '更新时间',   PRIMARY KEY (`id`) ) ENGINE=InnoDB AUTO_INCREMENT=17 DEFAULT CHARSET=utf8 COMMENT='合同表';",
            "sql": "SELECT year, budget_contract_amount, actual_contract_amount FROM contracts WHERE year = ? AND department = ?;",
            "sql_describe": "查询某个部门某年的年初预算的新签合同额和实际的新签合同额。"
        },
        {
            "ddl": "CREATE TABLE `income` (   `id` int(11) NOT NULL AUTO_INCREMENT COMMENT '收入ID',   `year` int(11) NOT NULL COMMENT '收入年份',   `department` varchar(255) NOT NULL COMMENT '部门名称',   `budget_contract_amount` decimal(10,2) NOT NULL COMMENT '年初预算的新签合同额',   `operating_income` decimal(10,2) NOT NULL COMMENT '营业收入',   `net_operating_income` decimal(10,2) NOT NULL COMMENT '营业净收入',   `cash_income` decimal(10,2) NOT NULL COMMENT '现金收入',   `net_cash_income` decimal(10,2) NOT NULL COMMENT '现金净收入',   `created_at` timestamp NOT NULL DEFAULT current_timestamp()  COMMENT '创建时间',   `updated_at` timestamp NOT NULL DEFAULT current_timestamp() ON UPDATE current_timestamp()  COMMENT '更新时间',   PRIMARY KEY (`id`) ) ENGINE=InnoDB AUTO_INCREMENT=17 DEFAULT CHARSET=utf8 COMMENT='收入指标表';",
            "sql": "SELECT year, budget_contract_amount, operating_income, net_operating_income, cash_income, net_cash_income FROM income WHERE year = ? AND department = ?;",
            "sql_describe": "查询某年某部门年初预算的新签合同额、营业收入、营业净收入、现金收入和现金净收入。"
        },
        {
            "ddl": "CREATE TABLE `income` (   `id` int(11) NOT NULL AUTO_INCREMENT COMMENT '收入ID',   `year` int(11) NOT NULL COMMENT '收入年份',   `department` varchar(255) NOT NULL COMMENT '部门名称',   `budget_contract_amount` decimal(10,2) NOT NULL COMMENT '年初预算的新签合同额',   `operating_income` decimal(10,2) NOT NULL COMMENT '营业收入',   `net_operating_income` decimal(10,2) NOT NULL COMMENT '营业净收入',   `cash_income` decimal(10,2) NOT NULL COMMENT '现金收入',   `net_cash_income` decimal(10,2) NOT NULL COMMENT '现金净收入',   `created_at` timestamp NOT NULL DEFAULT current_timestamp()  COMMENT '创建时间',   `updated_at` timestamp NOT NULL DEFAULT current_timestamp() ON UPDATE current_timestamp()  COMMENT '更新时间',   PRIMARY KEY (`id`) ) ENGINE=InnoDB AUTO_INCREMENT=17 DEFAULT CHARSET=utf8 COMMENT='收入指标表';",
            "sql": "SELECT SUM(net_cash_income) FROM income i WHERE i.`year` = ?",
            "sql_describe": "查询所有部门某年的现金净收入总和"
        },
        {
            "ddl": "CREATE TABLE `income` (   `id` int(11) NOT NULL AUTO_INCREMENT COMMENT '收入ID',   `year` int(11) NOT NULL COMMENT '收入年份',   `department` varchar(255) NOT NULL COMMENT '部门名称',   `budget_contract_amount` decimal(10,2) NOT NULL COMMENT '年初预算的新签合同额',   `operating_income` decimal(10,2) NOT NULL COMMENT '营业收入',   `net_operating_income` decimal(10,2) NOT NULL COMMENT '营业净收入',   `cash_income` decimal(10,2) NOT NULL COMMENT '现金收入',   `net_cash_income` decimal(10,2) NOT NULL COMMENT '现金净收入',   `created_at` timestamp NOT NULL DEFAULT current_timestamp()  COMMENT '创建时间',   `updated_at` timestamp NOT NULL DEFAULT current_timestamp() ON UPDATE current_timestamp()  COMMENT '更新时间',   PRIMARY KEY (`id`) ) ENGINE=InnoDB AUTO_INCREMENT=17 DEFAULT CHARSET=utf8 COMMENT='收入指标表';",
            "sql": "SELECT i.department,        SUM(i.operating_income) / NULLIF(                                           (SELECT SUM(operating_income)                                            FROM income                                            WHERE year = ?), 0) AS proportion FROM income i WHERE i.year = ? GROUP BY i.department;",
            "sql_describe": "查询某年每个部门营业收入在总营收中的占比"
        },
        {
            "ddl": "CREATE TABLE `contracts` (   `id` int(11) NOT NULL AUTO_INCREMENT COMMENT '合同ID',   `year` int(11) NOT NULL COMMENT '合同签订年份',   `department` varchar(255) NOT NULL COMMENT '部门名称,合同的所属部门',   `budget_contract_amount` decimal(10,2) NOT NULL COMMENT '年初预算的新签合同额',   `actual_contract_amount` decimal(10,2) NOT NULL COMMENT '实际的新签合同额',   `created_at` timestamp NOT NULL DEFAULT current_timestamp() COMMENT '创建时间',   `updated_at` timestamp NOT NULL DEFAULT current_timestamp() ON UPDATE current_timestamp() COMMENT '更新时间',   PRIMARY KEY (`id`) ) ENGINE=InnoDB AUTO_INCREMENT=17 DEFAULT CHARSET=utf8 COMMENT='合同表';  CREATE TABLE `income` (   `id` int(11) NOT NULL AUTO_INCREMENT COMMENT '收入ID',   `year` int(11) NOT NULL COMMENT '收入年份',   `department` varchar(255) NOT NULL COMMENT '部门名称',   `budget_contract_amount` decimal(10,2) NOT NULL COMMENT '年初预算的新签合同额',   `operating_income` decimal(10,2) NOT NULL COMMENT '营业收入',   `net_operating_income` decimal(10,2) NOT NULL COMMENT '营业净收入',   `cash_income` decimal(10,2) NOT NULL COMMENT '现金收入',   `net_cash_income` decimal(10,2) NOT NULL COMMENT '现金净收入',   `created_at` timestamp NOT NULL DEFAULT current_timestamp()  COMMENT '创建时间',   `updated_at` timestamp NOT NULL DEFAULT current_timestamp() ON UPDATE current_timestamp()  COMMENT '更新时间',   PRIMARY KEY (`id`) ) ENGINE=InnoDB AUTO_INCREMENT=17 DEFAULT CHARSET=utf8 COMMENT='收入指标表';",
            "sql": "SELECT * from (SELECT c.department as department,        SUM(c.actual_contract_amount) AS total_amount,        '实际合同金额最多部门' as mark FROM contracts c WHERE c.year = ? GROUP BY c.department ORDER BY total_amount DESC LIMIT 1) a union all SELECT * from (SELECT i.department as department,        SUM(i.operating_income) AS total_amount,        '营业收入金额最多部门' as mark FROM income i WHERE i.year = ? GROUP BY i.department ORDER BY total_amount DESC LIMIT 1) b",
            "sql_describe": "查询某年实际签订合同金额和营业收入金额最大的两个部门"
        }
    ]
    # 创建数据表与向量数据
    # db.create_milvus(data)

    # 检索条件向量化
    # embeddings = db.model.encode("查询条件")
    # print(embeddings.tolist()) # 查看所有向量
    # db.search_milvus(embeddings)
